using LinearAlgebra
using Plots, LaTeXStrings
pyplot()

Sigma_true = [
    2.  0.  0.  0.  0.  0.  0.  0;
    0.  2. -1.  0.  0.  0.  0.  0;
    0. -1.  2.  0. -1.  0.  0. -1;
    0.  0.  0.  2.  1.  0.  0.  0;
    0.  0. -1.  1.  2.  0.  0.  1;
    0.  0.  0.  0.  0.  2.  0.  0;
    0.  0.  0.  0.  0.  0.  2.  0;
    0.  0. -1.  0.  1.  0.  0.  2]

X = [-0.3311197  -0.3932145   0.7742677  -0.3097049  -1.5111846   0.3311429   2.1538895   0.8835870;
-2.7057869  -0.2514156  -0.6271286   0.5463020   1.6978879  -0.3421863  -0.6586413  -0.8402505;
-0.3192959   1.2968608  -0.4446002   1.2538625   1.8819350  -2.0727404  -1.9972990   0.3265276;
-0.4125173   0.5194493  -1.4735353   1.5864925   0.2298139   0.8494314   0.5313172  -0.7134051;
-2.7713921  -0.0423820  -1.1966070  -2.0107553  -0.2221868  -0.2953777  -1.7265336  -0.0383220;
-1.0180134  -0.8166305   0.4122348   0.0026914  -0.9920446   2.0909459  -0.4258249  -1.0972184;
 0.8650409   1.3432579  -0.6379710  -1.6605675  -2.1753728   0.9573123  -0.5445886   0.4765308;
-0.2625614   0.4310924  -1.1176806  -0.7204528   1.6350725   0.6776547   1.3796291   0.5220660;
-0.9123371  -1.9296817   1.3247186   1.4169735  -0.8366834  -0.5114306   1.4192098  -0.6890962;
-0.4228602   1.1161097   1.6705814   3.2483859   1.6894376  -0.1231031   1.1623457  -2.0525432;
-0.7095916  -1.2685852   0.1599799  -0.4224032   0.0252021   1.1433826  -0.7329448   1.9290979;
-0.4633842  -0.0797081   0.0926444   0.2468722   1.6129349   0.9928530   1.3698709   0.4049546;
-0.5948823  -0.3180217   0.6807469  -0.1785293   0.2297931   2.0016365  -0.3317563  -0.4840257;
-0.1052824  -2.3678404  -0.6171589  -1.5935395   1.1229458  -0.3642313   0.2468903   0.6588132;
-1.6527603  -2.1323678   1.6783155   3.0078456   0.7259103   0.0490899   0.4264523   0.4042916;
 3.0977721  -0.8027199   0.2392123  -2.8330579  -0.9121670  -0.8299413  -1.9825168  -0.1996475;
 0.6697608   1.0988143  -1.0885298   0.8310039  -0.6695225   1.5023243   0.0969627  -1.8685160;
 1.8485790  -0.4324271  -1.7509966  -0.9744788   1.2262814  -0.3216769  -1.7407050   0.5857962;
 0.4199998   0.1054119   0.7350115   1.7864235  -0.0520600  -0.7381368  -1.8674077  -1.1380628;
 0.6701009  -0.5945656  -1.3242998   0.1592060   1.1155564  -2.6386036  -1.0102455   0.2825183;
 0.5837604  -1.4277136  -1.0824538   0.9759325   0.6190599  -1.3624196  -1.1606201   2.3861869;
 0.4824618   1.4701000  -0.9503183   0.3830496  -0.5714352   0.1090383  -1.1537300  -0.3257069;
-1.2122038  -1.6077674   0.3301125   1.4432326   1.0597755  -0.8848272   3.8469001  -0.8128089;
-2.1423175  -1.2859498  -0.3093076   0.4973710   0.5315224   1.1975472  -0.1029956   0.3706120;
 0.2174001   1.2404896  -0.0527978  -1.7095400  -0.5960344   1.2499632   0.5449211  -0.0448455;
-1.3262849  -1.5899511   0.7485016  -0.2622975  -0.4558745   1.9483187   1.4903012   0.8006743;
 0.9212063  -4.7382379   2.5271192   3.1098470   1.5434335  -1.3491622   1.6059323  -0.2300004;
-2.6407118   1.5776585  -0.2052617   0.3756392   1.1194163  -3.2734215   0.1157875   0.7644046;
 0.9610964  -1.1207784  -0.2378744   1.0983257   1.3869479   1.0115922  -1.5241613   0.1350005;
 0.2636862  -1.5156891  -0.8902854  -2.3923297  -0.5273719   2.8639841   0.8964969   1.3435503;
-0.4730562   0.1234838  -0.9236915  -0.4000590   1.2714111   1.6974755   1.2243627   0.9603859;
 1.4601264  -0.5659824   1.9538219  -0.0001947  -1.1066020  -3.0011821   1.2980528  -3.6546994;
-0.0841816  -0.9661943   0.2898518  -2.8143836  -0.3795513   1.3108776   1.0954980  -2.6934468;
 0.1848951   0.7326472  -0.0012979  -2.3559473  -0.4322298   0.6223169  -2.0225253   2.3752747;
-0.8488349   0.3068909   1.5376214   0.1497557  -1.2388319  -0.2815129   0.6535130  -1.9475487;
-2.7606690  -2.1736976   0.2587605  -0.6145264   0.1003725  -3.0165647  -1.7505355   1.0967079;
 1.2015188   0.8893968   0.3391192  -0.6410825  -0.5730524   0.8334876  -0.6972086  -0.2614790;
-0.7922156  -1.3573571   1.2776654  -0.4787305  -1.3617914   1.1723760   0.8904296   0.1550097]

lam = .15
N = size(X)[1]
Sigma_emp = X'*X./N

function plot_helper(Sigma_est; save_dest="")
    names = [L"$\Sigma_{true}$", L"$\hat{\Sigma}$", L"$\Sigma_{emp}$"]
    mats = [Sigma_true, Sigma_est, Sigma_emp]
    plot_array = [] 
    hmin = minimum([minimum(mats[i]) for i in 1:3])
    hmax = maximum([maximum(mats[i]) for i in 1:3])
    for i in 1:3
        name = names[i]
        mat = mats[i]
        push!(plot_array, heatmap(mat,yflip=true,title=name,aspect_ratio=:equal,showaxis=false,grid=false, clims=(hmin, hmax)))
    end
    p = plot(plot_array...,layout = (1,3))
    if save_dest == ""
        display(p)
        return p
    else
        savefig(p, save_dest)
    end
end
